#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <assert.h>
#include "dnsclient.h"

void
free_addrinfo(struct dc_addrinfo **getaddrinfo)
{
    struct dc_addrinfo *addrinfo;

    assert(getaddrinfo);

    addrinfo = *getaddrinfo;

    if (addrinfo == NULL)
	return;

    free_addrinfo(&(addrinfo->ai_next));

    free(addrinfo->ai_addr);
    free(addrinfo);

    *getaddrinfo = addrinfo;
}

void
add_addrinfo(struct dc_addrinfo **getaddrinfo, uint32_t *s_addr, int family)
{
    struct dc_addrinfo *addrinfo;

    assert(getaddrinfo);

    addrinfo = *getaddrinfo;
    
    if (addrinfo == NULL) {
	addrinfo = malloc(sizeof(struct dc_addrinfo));
	
	if (family == AF_INET) {
	    struct sockaddr_in *sin;
    	    sin = malloc(sizeof(struct sockaddr_in));
    	    memset (sin, 0, sizeof(struct sockaddr_in));
    	    sin->sin_family = AF_INET;
    	    sin->sin_addr.s_addr = *s_addr;
    
	    addrinfo->ai_family = AF_INET;
	    addrinfo->ai_addrlen = sizeof(struct sockaddr_in);
    	    addrinfo->ai_addr = ((struct sockaddr *) sin);
	} else if (family == AF_INET6) {
	    struct sockaddr_in6 *sin6;
    	    sin6 = malloc(sizeof(struct sockaddr_in6));
    	    memset (sin6, 0, sizeof(struct sockaddr_in6));
    	    sin6->sin6_family = AF_INET6;
	    int i;
	    for (i = 0; i < 4; i++)
    		sin6->sin6_addr.s6_addr32[i] = s_addr[i];

	    addrinfo->ai_family = AF_INET6;
	    addrinfo->ai_addrlen = sizeof(*sin6);
    	    addrinfo->ai_addr = ((struct sockaddr *) sin6);
	}

	addrinfo->ai_next = NULL;
	*getaddrinfo = addrinfo;
    } else {
	add_addrinfo(&(addrinfo->ai_next), s_addr, family);
    }
}

size_t
dns_name_to_label (const char *name, unsigned char *label, size_t size)
{
    const char *p, *q;
    unsigned char length, *rname;

    if (strlen(name) + 1 >= (size -1))
	return 0;

    memset (label, 0, size);

    p = name;
    rname = label;

    while (p && (*p != '\0')) {
	int more;
	q = strchr (p, '.');
	if (q == NULL) {
	    more = 0;
	    length = strlen(p);
	} else {
	    more = 1;
	    length = q - p;
	}

	*rname++ = length;
	memmove (rname, p, length);
	rname += length;

	p = q + more;
    }

    *rname++ = 0;

    return (rname - label);
}

void
dns_format_header(struct dns_query_header *header)
{
    /* set dns packet's header part */
    memset (header, 0, sizeof(struct dns_query_header));

    /* only sets the values those required for query because we already set all bits to 0 */
    header->dns_id	= (int) getpid();	// get process id as dns_id
    header->dns_rd	= 1;			// rescursion shold be on
    header->dns_qr	= 0;			// It's Query
    header->dns_opcode	= 0;			// It's standard query
    header->dns_qdcount	= 1;			// Only one query to process

    /* converting all the values to network byte order */
    header->dns_id	= htons(header->dns_id);
    header->dns_qdcount	= htons(header->dns_qdcount);
    header->dns_ancount	= htons(header->dns_ancount);
    header->dns_nscount	= htons(header->dns_nscount);
    header->dns_arcount	= htons(header->dns_arcount);
}

size_t
dns_format_query (const char *host, dns_class_t class, dns_type_t type, unsigned char *buffer, size_t buffer_len)
{
    size_t label_length = 0;
    unsigned char *t_buffer;
    struct dns_query_header header;
    unsigned char label[buffer_len];

    t_buffer = buffer;

    dns_format_header(&header);

    /* copy header in query buffer */
    memcpy (t_buffer, &header, sizeof(header));
    t_buffer += sizeof(header);

    /* convert query domain name in label format and get its label length */
    label_length = dns_name_to_label (host, label, sizeof(label));
    if (label_length == 0)
	return 0;

    /* append label in dns header */
    memcpy (t_buffer, label, label_length); 
    t_buffer += label_length;

    /* type and class are having u_int16_t data type and requires 2 byte so first byte is set to 0 for both */
    *((u_int16_t *) t_buffer) = ntohs(type);
    t_buffer += sizeof(u_int16_t);
    *((u_int16_t *) t_buffer) = ntohs(class);
    t_buffer += sizeof(u_int16_t);

    /* returning the original size of query */
    return (t_buffer - buffer);
}

static size_t
parse_qdcount_results()
{
}

static size_t
parse_header_results(const unsigned char *results, struct dns_query_header *header)
{
    /* copy header from the results */
    memcpy (header, results, sizeof(struct dns_query_header));
    header->dns_id = ntohs(header->dns_id);
    header->dns_qdcount = ntohs(header->dns_qdcount);
    header->dns_ancount = ntohs(header->dns_ancount);
    header->dns_nscount = ntohs(header->dns_nscount);
    header->dns_arcount = ntohs(header->dns_arcount);

    /* query response should be answer */
    if (header->dns_qr != 1)
	return 0;

    /* there should not be error in query */
    if (header->dns_rcode != 0)
	return 0;

    return sizeof(*header);
}

struct dc_addrinfo *
dns_parse_results (const unsigned char *results, size_t length)
{
    int i, j;

    size_t ret_size = 0, label_size = 0;

    unsigned char buffer[length];
    const unsigned char *p = NULL;

    struct dns_rr *res = NULL;
    struct dns_query_header header;

    struct dc_addrinfo *addrinfo = NULL;

    /* porcess header from dns answer */
    ret_size = parse_header_results(results, &header);
    if (ret_size == 0)
	return NULL;
    p = results + sizeof(header);

    /* allocate memory for all returned records */
    res = calloc (header.dns_qdcount + header.dns_ancount + header.dns_nscount + header.dns_arcount, sizeof(struct dns_rr));
    if (res == NULL)
	return NULL;

    /* process query section from the dns answer */
    for (i = 0; i < header.dns_qdcount; i++) {
	/* parse label got in dns query section of the dns answer  */
	label_size = dns_parse_label (results, length - (p - results), p, buffer, sizeof(buffer));
	if (label_size == 0) {
	    free(res);
	    return NULL;
	}
	
	p += label_size;

	/* question section contains type and class as well, so filling it.*/
	res[i].dns_type = ntohs(*(int16_t*)p);
	p += sizeof(u_int16_t);
	res[i].dns_class = ntohs(*(int16_t*)p);
	p += sizeof(u_int16_t);

#ifdef DEBUG
	fprintf(stderr, "Queried for '%s', type = %d, class = %d\n", buffer, res[i].dns_type, res[i].dns_class);
#endif
    }

    /* p is now at answer section of resource records, process records from answer section */
    for (i = 0; i < header.dns_ancount + header.dns_nscount + header.dns_arcount; i++) {
	char *tmp;

#ifdef DEBUG
    	char addr[INET_ADDRSTRLEN], addr6[INET6_ADDRSTRLEN];
#endif

	label_size = dns_parse_label (results, length - (p - results), p, buffer, sizeof(buffer));
	if (label_size == 0) {
	    free(res);
	    return NULL;
	}
	
	/* copy parsed label in tmp and reserve (header + label) section and process reamining part */
	tmp = strdup((char *)buffer);
	p += label_size;

	res[i].dns_type = ntohs(*(u_int16_t*)p);
	p += sizeof(u_int16_t);

	res[i].dns_class = ntohs(*(u_int16_t*)p);
	p += sizeof(u_int16_t);

	res[i].dns_ttl = ntohl(*(int32_t*)p);
	p += sizeof(int32_t);

	res[i].dns_rlength = ntohs(*(u_int16_t*)p);
	p += sizeof(u_int16_t);

	/* hear we are getting resource record */
	switch (res[i].dns_type) {
	    case DNS_T_CNAME:
		label_size = dns_parse_label (results, length - (p - results), p, buffer, sizeof(buffer));
		if (label_size == 0) {
		    free(res);
		    return NULL;
		}

#ifdef DEBUG
		res[i].dns_rdata.cname.name = strdup((char *)buffer);
		fprintf(stderr, "Answer for '%s', CNAME : '%s', type = %d, class = %d, ttl = %d, rlength = %d\n", tmp, res[i].dns_rdata.cname.name, res[i].dns_type, res[i].dns_class, res[i].dns_ttl, res[i].dns_rlength);
		free((void *)res[i].dns_rdata.cname.name);
#endif

		p += res[i].dns_rlength;
		

		break;
	    case DNS_T_A:
		/* inet_ntop API requires address in network byte order so no need to convert it in host bye order */
		res[i].dns_rdata.a.address = *(uint32_t*)p;

		add_addrinfo(&addrinfo, &(res[i].dns_rdata.a.address), AF_INET);

#ifdef DEBUG
		memset (addr, '\0', sizeof(addr));
		inet_ntop(AF_INET, &(res[i].dns_rdata.a.address), addr, sizeof(addr));
		fprintf(stderr, "Answer for '%s', Address : '%s', type = %d, class = %d, ttl = %d, rlength = %d\n", tmp, addr, res[i].dns_type, res[i].dns_class, res[i].dns_ttl, res[i].dns_rlength);
#endif

		p += res[i].dns_rlength;

		break;
	    case DNS_T_AAAA:
		/* inet_ntop API requires address in network byte order so no need to convert it in host bye order */
		for (j = 0; j < 4; j++) {
    		    res[i].dns_rdata.aaaa.oct[j] = *(uint32_t*)(p + (4 * j));
		}

		add_addrinfo(&addrinfo, res[i].dns_rdata.aaaa.oct, AF_INET6);

#ifdef DEBUG
		memset (addr6, '\0', sizeof(addr6));
		inet_ntop(AF_INET6, &(res[i].dns_rdata.aaaa.oct), addr6, sizeof(addr6));
		fprintf(stderr, "Answer for '%s', Address : '%s', type = %d, class = %d, ttl = %d, rlength = %d\n", tmp, addr6, res[i].dns_type, res[i].dns_class, res[i].dns_ttl, res[i].dns_rlength);
#endif

		p += res[i].dns_rlength;

		break;
	    default:
		label_size = dns_parse_label (results, length - (p - results), p, buffer, sizeof(buffer));

		if (label_size == 0) {
		    free(res);
		    return NULL;
		}
		p += res[i].dns_rlength;

		break;
	}
	free(tmp);
    }

    free(res);
    return addrinfo;
}

size_t
dns_parse_label (const unsigned char *base, size_t base_length, const unsigned char *label, unsigned char *output, size_t output_length)
{
    size_t ret = 0;
    int pointer = 1;

    const unsigned char *p;

    if (label == NULL)
	return 0;

    p = label;

    memset (output, '\0', output_length);

    /* address of passed label mest be less than answer's last address */
    // while (p && (p < (base + base_length) && *p)) {
    while (p && *p) {
	/* This is for offset calculation when response comes compressed */
	/* If (*p >= 192) it means its an offset of 2 bytes and value at the next address should be called */
	if (*p >= 0xC0) {
	    p = base + (ntohs(*(u_int16_t *)p) & 0xC00000FF);
	    if (pointer)
    		ret += 2;
	    pointer = 0;
	    continue;
	}

	/* This is number that gives no of chars to read */
	if (pointer)
	    ret += (*p + 1);

	strncat((char *)output, (char *)p + 1, *p);
	strcat((char *)output, ".");
	/* Now increase p to ret next (*p +1) bytes */
	p += (*p +1);
    }

    return pointer ? ret + 1 : ret;
}

struct dc_addrinfo *
dns_client_query (const char *host, u_int16_t rr_type, int udpsock)
{
    dns_class_t dns_class = DEFAULT_DNS_CLASS;

    unsigned char *qbuffer, *abuffer;
    // struct dns_rr *results = NULL;
    struct dc_addrinfo *results = NULL;

    ssize_t qsize, asize, ret = 0;

    qsize = DEFAULT_QSIZE;
    qbuffer = malloc(qsize);
    if (qbuffer == NULL)
	return NULL;

    do {
	ret = dns_format_query (host, dns_class, rr_type, qbuffer, qsize);
	if (ret > 0) {
	    qsize = ret;
	    break;
	}
	free(qbuffer);
	qsize += DEFAULT_QSIZE;
	qbuffer = malloc(qsize);
	if (qbuffer == NULL)
	    return NULL;
    } while (qsize <= MAX_QSIZE);

    if ((ret > 0) && (qsize >0)) {
	asize = DEFAULT_ASIZE;
	abuffer = malloc(asize);

	/* send generated dns query qbuffer to server */	
	ret = write (udpsock, qbuffer, qsize);
	if (ret == -1) {
	    fprintf(stderr, "ERROR: fail to write on server\n");
	    exit (1);
	}

	/* receive dns query answer in abuffer */
	memset(abuffer, '\0', sizeof(asize));
	ret = read (udpsock, abuffer, asize);
	if (ret == -1) {
	    fprintf(stderr, "Failed to read server response\n");
	    exit(1);
	} else if (ret == 0)
	    fprintf(stderr, "Connection cosed by peer\n");
	/* parse received results in user readable format */
	if (ret > 0)
	    results = dns_parse_results (abuffer, ret);

	/* free all allocated buffer */
	free(abuffer);
	free(qbuffer);
    }

    return results;
}
